package cn.iocoder.yudao.module.engine.service.core.queue.llm;


import cn.iocoder.yudao.module.infra.llm.model.LLMMessage;

import java.util.*;
import java.util.concurrent.locks.ReentrantReadWriteLock;

public class LLMMessageQueue {
    private final LinkedHashMap<String, LLMMessage> messageMap;
    private final int maxSize;
    private final ReentrantReadWriteLock lock = new ReentrantReadWriteLock();

    public LLMMessageQueue(int maxSize) {
        if (maxSize <= 0) {
            throw new IllegalArgumentException("Max size must be positive");
        }
        this.maxSize = maxSize;
        this.messageMap = new LinkedHashMap<String, LLMMessage>(maxSize, 0.75f, false) {
            @Override
            protected boolean removeEldestEntry(Map.Entry<String, LLMMessage> eldest) {
                return size() > LLMMessageQueue.this.maxSize;
            }
        };
    }

    /**
     * 添加新消息到队列
     * @param message 要添加的消息
     * @return 如果messageId已存在返回false，添加成功返回true
     * @throws IllegalArgumentException 如果消息或messageId为null
     */
    public boolean add(LLMMessage message) {
        if (message == null || message.getMessageId() == null) {
            throw new IllegalArgumentException("Message and messageId cannot be null");
        }

        lock.writeLock().lock();
        try {
            // 如果已存在相同messageId的消息，返回false
            if (messageMap.containsKey(message.getMessageId())) {
                return false;
            }

            messageMap.put(message.getMessageId(), message);
            return true;
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * 获取队列中的所有消息列表（按插入顺序）
     * @return 消息列表的不可修改视图
     */
    public List<LLMMessage> getMessageList() {
        lock.readLock().lock();
        try {
            return Collections.unmodifiableList(new ArrayList<>(messageMap.values()));
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * 根据messageId获取消息
     * @param messageId 消息ID
     * @return 对应的消息，如果不存在返回null
     */
    public LLMMessage get(String messageId) {
        lock.readLock().lock();
        try {
            return messageMap.get(messageId);
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * 根据messageId移除消息
     * @param messageId 要移除的消息ID
     * @return 被移除的消息，如果不存在返回null
     */
    public LLMMessage remove(String messageId) {
        lock.writeLock().lock();
        try {
            return messageMap.remove(messageId);
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * 清空队列
     */
    public void clear() {
        lock.writeLock().lock();
        try {
            messageMap.clear();
        } finally {
            lock.writeLock().unlock();
        }
    }

    /**
     * 获取当前队列大小
     * @return 队列中的消息数量
     */
    public int size() {
        lock.readLock().lock();
        try {
            return messageMap.size();
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * 检查队列是否为空
     * @return 如果队列为空返回true，否则返回false
     */
    public boolean isEmpty() {
        lock.readLock().lock();
        try {
            return messageMap.isEmpty();
        } finally {
            lock.readLock().unlock();
        }
    }

    /**
     * 检查是否包含指定messageId的消息
     * @param messageId 要检查的消息ID
     * @return 如果包含返回true，否则返回false
     */
    public boolean contains(String messageId) {
        lock.readLock().lock();
        try {
            return messageMap.containsKey(messageId);
        } finally {
            lock.readLock().unlock();
        }
    }
}